/**
 *  @file SSLSocketFactory.cpp
 */
#include "SSLSocketFactory.h"
#include "../../common/CppMemPool.h"
#include "../../common/UtilFuncT.h"
#include "../../collections/KeyValueListT.h"
#include "../Socket.h"
#include "SSLContext.h"
#include "SSLSocket.h"
#include "SSLServerSocket.h"
#include "SSLException.h"
#include "SSLSocketException.h"
#include "../SocketException.h"
#include "../../security/cert/X509Certificate.h"
#include "../../convert/Convert.h"
#include "../../common/ObjPtrT.h"

using namespace cppflib::security::cert;
using namespace cppflib::collections;
using namespace cppflib::exceptions;

namespace cppflib
{

namespace net
{

namespace ssl
{

/**
 *  An internal callback to verify target host name connecting to against that exists in peer cert
 */
class VerifyHostNameCallback : public ISSLPeerCertCallback
{
private:
   CFString targetHostName;
   SSLSocketFactory::verifyhostname_t verifyType;
   ISSLPeerCertCallback * pOrgCallback;
   
private:
   VerifyHostNameCallback() { }

   /**
    *  @note host names are case-insensitive
    *
    *  @param[in] hostName -- may be a pattern (e.g. *.google.com) or a normal host name (e.g. www.google.com)
    *                         but *CANNOT* contain double wildcard (e.g. *.*.abc.com)
    *
    *  @return true if hostname matches the pre-set targetHostName
    */
   bool CheckHostName(const CFString & hostNamePattern)
   {
      if (this->targetHostName.IsEmpty())
         return false;
     
      return IsHostNameMatched(targetHostName, hostNamePattern);
   }

   /**
    *  @return true if target hostname match that in subjectAltName -> DNS
    */
   bool CheckSubjectAltName(X509Certificate &x509Cert)
   {
      for (int i = 0; i < x509Cert.GetExtensionCount(); ++i) {
         CFString name;
         if (!x509Cert.GetExtensionName(i, name))
            break;

         if (name != _S("subjectAltName"))
            continue;

         X509Certificate::extensiontype_t type = security::cert::X509Certificate::EXT_UNKNOWN;
         if (!x509Cert.GetExtensionType(i, type))
            break;

         if (type != security::cert::X509Certificate::EXT_LIST) // subjectAltName should be a multi-value list
            continue;

         NameValueList extValList;
         if (!x509Cert.GetExtensionValueList(i, extValList))
            break;

         for (int j = 0; j < static_cast<int>(extValList.GetCount()); ++j) {
            if (*extValList.GetKey(j) == _S("DNS")) {    // check each 'DNS'
               if (CheckHostName(*extValList.GetVal(j)))
                  return true;
            }
         }
         
         break; // fail if all DNS don't match
      }

      return false;
   }

   /**
    *  @return true if target hostname match that in common name (CN)
    */
   bool CheckCommanName(X509Certificate &x509Cert)
   {
      CFString commonName;
      if (!x509Cert.GetOneValueFromSubjectName(CFString("CN"), commonName))
         return false;

      return CheckHostName(commonName);
   }

public:
   VerifyHostNameCallback(SSLSocketFactory::verifyhostname_t verifyType, ISSLPeerCertCallback * pOrgCallback)
   {
      this->verifyType = verifyType;
      this->pOrgCallback = pOrgCallback;

      if (this->pOrgCallback)
         cppmempool::Retain(this->pOrgCallback);
   }

   virtual ~VerifyHostNameCallback()
   {
      if (this->pOrgCallback)
         cppmempool::Release(this->pOrgCallback);
   }

   /**
    *  Do nothing, just pass on to original callback if any
    */
   virtual void ReadPeerCert(bool isCertOK, X509Certificate &x509Cert)
   {
      if (pOrgCallback)
         pOrgCallback->ReadPeerCert(isCertOK, x509Cert);
   }

   /**
    *  Verify target host name if required
    */
   virtual void ReadPeerCertAfterConnection(X509Certificate &x509Cert)
   {
      if (pOrgCallback)
         pOrgCallback->ReadPeerCertAfterConnection(x509Cert); // pass to original callback first

      if (verifyType == SSLSocketFactory::VERIFYHOSTNAME) {
         if (!CheckSubjectAltName(x509Cert) && !CheckCommanName(x509Cert))
            throw SSLException(_S("Host name not matched !"));
      }
   }

   /**
    *  Set current and original callback's host name
    */
   virtual void SetTargetHostName(pcwstr_t hostName)
   {
      this->targetHostName = hostName;
      if (pOrgCallback)
         pOrgCallback->SetTargetHostName(hostName);
   }

};

/*************** SSLSocketFactory ****************/

/**
 *  SSL Socket created using this constructor won't check server's cert, which is useful for testing purpose
 */
SSLSocketFactory::SSLSocketFactory(void)
{
   try {
      this->pSSLCtxt = new SSLContext();
      cppmempool::Retain(this->pSSLCtxt);
      this->pPeerCertCb = NULL;
      this->verifyType = SSLSocketFactory::NOTVERIFY;
   }
   catch(BaseException & e) {
      ReleaseBeforeThrowing();
      throw e;
   }
}

/**
 *  This client won't send any personal cert to server
 *
 *  @param [in] rootCertPath -- path to a PEM file with all trusted root CA
 *  @param [in] pPeerCertCb -- call back to check peer cert. Pass NULL if not interested
 */
SSLSocketFactory::SSLSocketFactory(const CFString &rootCertPath, ISSLPeerCertCallback *pPeerCertCb)
{
   try{
      this->pSSLCtxt = new SSLContext(false, rootCertPath, true);
      cppmempool::Retain(this->pSSLCtxt);
      this->pPeerCertCb = pPeerCertCb;
      cppmempool::Retain(this->pPeerCertCb);
      this->verifyType = SSLSocketFactory::VERIFYHOSTNAME; // verify hostname by default
   }
   catch(BaseException & e) {
      ReleaseBeforeThrowing();
      throw e;
   }
}

/**
 *  This client will send a personal cert to server
 *
 *  @param [in] rootCertPath -- path to a PEM file with all trusted root CA
 *  @param [in] pPeerCertCb -- call back to check peer cert. Pass NULL if not interested
 *  @param [in] clientCertPath -- path to client's personal cert in PEM format
 *  @param [in] pPrivateKeyPwdCb -- call back to get the password for client's personal cert
 */
SSLSocketFactory::SSLSocketFactory(const CFString &rootCertPath, ISSLPeerCertCallback *pPeerCertCb,
                                   const CFString &clientCertPath, ISSLPrivateKeyPwdCallback *pPrivateKeyPwdCb)
{
   try {
      this->pSSLCtxt = new SSLContext(false, rootCertPath, true);
      cppmempool::Retain(this->pSSLCtxt);
      this->pPeerCertCb = pPeerCertCb;
      cppmempool::Retain(this->pPeerCertCb);
      this->verifyType = SSLSocketFactory::VERIFYHOSTNAME; // verify hostname by default

      this->pSSLCtxt->SetCertChain(clientCertPath, pPrivateKeyPwdCb);
   }
   catch(BaseException & e) {
      ReleaseBeforeThrowing();
      throw e;
   }
   
}
    
SSLSocketFactory::~SSLSocketFactory(void)
{
   ReleaseBeforeThrowing();
}

void SSLSocketFactory::ReleaseBeforeThrowing()
{
   if (pPeerCertCb) {
      cppmempool::Release(pPeerCertCb);
      pPeerCertCb = NULL;
   }
   
   if (pSSLCtxt) {
      cppmempool::Release(pSSLCtxt);
      pSSLCtxt = NULL;
   }
}

/**
 *  @param [in] isSSL -- true to create a SSL client socket; false to create normal TCP socket
 */
Socket * SSLSocketFactory::CreateSocket(bool isSSL, const CFString &socketAddress, u16_t socketPort)
{
   if (!isSSL)
      return SocketFactory::CreateSocket(isSSL, socketAddress, socketPort);

   SSLSocket *pSock = pSSLCtxt->CreateClientSocket(NULL, new VerifyHostNameCallback(verifyType, pPeerCertCb));
   pSock->Create(socketAddress, socketPort);
   return pSock;
}

/**
 *  Not implemented for this class
 */
bool SSLSocketFactory::CreateListeningSockets(collections::ListCppObjT<Socket> &sockList, int backLog)
{
   throw SSLException(_S("SSLSocketFactory is not for creating listening sockets"));
   return false;
}

/**
 *  Set verification type
 */
void SSLSocketFactory::SetVerifyHostnameType(SSLSocketFactory::verifyhostname_t type)
{
   this->verifyType = type;
}

/**
 *  Get verification type
 */
SSLSocketFactory::verifyhostname_t SSLSocketFactory::GetVerifyHostnameType()
{
   return this->verifyType;
}

/*************** SSLVHost ****************/

/**
 *  This internal structure contains info about a SSL virtual host
 */
class SSLVHost : public BaseObject
{
public:
      CFString *hostAddr;        ///< should be human readable domain name
      CFString *serverCertPath;  ///< path to server's certificate (contains private key)
      ISSLPeerCertCallback *pPeerCertCb;  ///< callback when there is an incoming SSL client 
      ISSLPrivateKeyPwdCallback *pPrivateKeyPwdCb;  ///< callback to get password of private key of server's certificate
      bool needClientAuth;   ///< whether a cert should have been sent from client during SSL negotiation
      SSLContext *pSSLCtxt;  ///< context to create a listening server socket

private:
   SSLVHost() { }

public:
   SSLVHost(const CFString & hostAddr, const CFString & serverCertPath,
                     ISSLPrivateKeyPwdCallback *pPrivateKeyPwdCb, ISSLPeerCertCallback *pPeerCertCb,
                     bool needClientAuth)
   {
      this->hostAddr = new CFString(hostAddr);
      this->hostAddr->MakeUpper();  // make host name upper case
      cppmempool::Retain(this->hostAddr);
      this->serverCertPath = new CFString(serverCertPath);
      cppmempool::Retain(this->serverCertPath);
      this->pPrivateKeyPwdCb = pPrivateKeyPwdCb;
      cppmempool::Retain(this->pPrivateKeyPwdCb);
      this->pPeerCertCb = pPeerCertCb;
      cppmempool::Retain(this->pPeerCertCb);
      this->needClientAuth = needClientAuth;
      this->pSSLCtxt = NULL;
   }

   virtual ~SSLVHost()
   {
      if (hostAddr) cppmempool::Release(hostAddr);
      if (serverCertPath) cppmempool::Release(serverCertPath);
      if (pPeerCertCb) cppmempool::Release(pPeerCertCb);
      if (pPrivateKeyPwdCb) cppmempool::Release(pPrivateKeyPwdCb);
      if (pSSLCtxt) cppmempool::Release(pSSLCtxt);
   }

   void SetAttachingSSLContext(SSLContext *pSSLCtxt)
   {
      cppmempool::Retain(pSSLCtxt);
      cppmempool::Release(this->pSSLCtxt);
      this->pSSLCtxt = pSSLCtxt;
      
   }
};

/*************** Find matching hostname callback ***************/

class ServerFindMatchingHostNameCallback : public ISSLFindMatchingHostNameCallback
{
private:
   /*
    *  Make this weak so as to avoid circular reference, since
    *
    *     ServerFindMatchingHostNameCallback -------> pVirtualHostList
    *     pVirtualHostList -------------------------> SSLVHost
    *     SSLVHost ---------------------------------> SSLContext
    *     SSLContext -------------------------------> ServerFindMatchingHostNameCallback
    */
   WeakReferenceT<ListCppObjT<BaseObject> > * pVirtualHostList;

private:
   ServerFindMatchingHostNameCallback() { }

public:
   ServerFindMatchingHostNameCallback(ListCppObjT<BaseObject> * pVirtualHostList)
   {
      this->pVirtualHostList = new WeakReferenceT<ListCppObjT<BaseObject> >(pVirtualHostList);
      cppmempool::Retain(this->pVirtualHostList);
   }

   virtual ~ServerFindMatchingHostNameCallback()
   {
      cppmempool::Release(this->pVirtualHostList);
   }

   virtual SSLContext *GetMatchingSSLContext(const CFString &hostName)
   {
      ListCppObjT<BaseObject> * pVHList = pVirtualHostList->Ptr();

      if (pVHList) {
         for (int i = 0; i < static_cast<int>(pVHList->GetCount()); ++i) {
            BaseObject *pObj = pVHList->Get(i);
            if (!cppflib::InstanceOf<SSLVHost>(pObj))
               continue;

            SSLVHost * pSingleHost = static_cast<SSLVHost*>(pObj);
            if (hostName.CompareNoCase(*pSingleHost->hostAddr) == 0)
               return pSingleHost->pSSLCtxt;
         }
      }

      return NULL; // not found
   }

};


/*************** SSLVirtualHostInfo ****************/

SSLVirtualHostInfo::SSLVirtualHostInfo()
{
}

/**
 *  To setup the first (default) virtual host info
 *
 *
 *  SSLVirtualHostInfo is structured like this:
 *
 *
 *  SSLVirtualHostInfo ---+-- listenPort (shared among virtual hosts listening on the same port)
 *                        |
 *                        +-- SSLVHost1 --+-- hostAddr
 *                        |               |
 *                        |               +-- serverCertPath
 *                        |               |
 *                        |               +-- ...
 *                        |               |
 *                        |               +-- pSSLCtxt
 *                        |
 *                        |
 *                        +-- SSLVHost2 --+-- hostAddr
 *                                        |
 *                                        +-- serverCertPath
 *                                        |
 *                                        +-- ...
 *                                        |
 *                                        +-- pSSLCtxt
 *
 *  @param [in] listenPort -- the port to listen
 *  @param [in] bindAddr -- a specific IP address or empty string to listen on all interfaces
 *  @param [in] hostAddr -- human readable domain name (default virtual host)
 *
 */
SSLVirtualHostInfo::SSLVirtualHostInfo(u16_t listenPort, const CFString & bindAddr, const CFString & hostAddr, 
                                 const CFString & serverCertPath,
                                 ISSLPrivateKeyPwdCallback *pPrivateKeyPwdCb, ISSLPeerCertCallback *pPeerCertCb,
                                 bool needClientAuth) 
                            : VirtualHostInfo(listenPort, bindAddr, CFString::Empty())
{
   //
   // For VirtualHostInfo (base class), pVirtualHostList holds list of host names only. 
   // However for SSLVirtualHostInfo, pVirtualHostList holds the 'SSLVHost' structure.
   //
   pVirtualHostList->AddTail(new SSLVHost(hostAddr, serverCertPath, pPrivateKeyPwdCb,
                                                   pPeerCertCb, needClientAuth));
}

SSLVirtualHostInfo::~SSLVirtualHostInfo()
{
   if (pVirtualHostList) {
      cppmempool::Release(pVirtualHostList);
      pVirtualHostList = NULL;
   }
}

/** 
 *  This one is not supported here
 */
bool SSLVirtualHostInfo::AddVirtualHost(const CFString & hostAddr)
{
   throw SSLException(_S("SSL virtual host must be specified with a server certificate"));
   return false;
}

/**
 *  Add one more virtual host other than the default which also binds to the same [address, port] pair
 *
 *  @hostAddr -- should be a human readable domain name but not IP address
 */
bool SSLVirtualHostInfo::AddVirtualHost(const CFString & hostAddr, const CFString & serverCertPath,
                                     ISSLPrivateKeyPwdCallback *pPrivateKeyPwdCb, ISSLPeerCertCallback *pPeerCertCb,
                                     bool needClientAuth)
{
   for (int i = 0 ; i < static_cast<int>(pVirtualHostList->GetCount()); ++i) {
      SSLVHost * pSSLVHost = static_cast<SSLVHost*>(pVirtualHostList->Get(i));
      if (hostAddr.CompareNoCase(*pSSLVHost->hostAddr) == 0) // host name must be different
         return false;
   }

   pVirtualHostList->AddTail(new SSLVHost(hostAddr, serverCertPath, pPrivateKeyPwdCb,
                                                   pPeerCertCb, needClientAuth));
   return true;
}


/*************** SSLServerSocketFactory ****************/

SSLServerSocketFactory::SSLServerSocketFactory(void)
{
   this->pVirtualHostInfoList = NULL;
}


PRIVATE const int DEFAULT_SSLCONTEXT_INDEX = 0;

/**
 *
 *  @param [in] rootCertPath -- path to a PEM file with all trusted root CA for verifying peer
 *  @param [in] pVirtualHostInfoList -- virtual host info list. Elements in this list may be a 
 *                                      'VirtualHostInfo' or a 'SSLVirtualHostInfo'
 */
SSLServerSocketFactory::SSLServerSocketFactory(const CFString &rootCertPath,
                                               ListCppObjT<VirtualHostInfo> * pVHInfoList)
                                              : SocketFactory(pVHInfoList)
{
      
   
   for (int i = 0; i < static_cast<int>(this->pVirtualHostInfoList->GetCount()); ++i) {
      if (!cppflib::InstanceOf<SSLVirtualHostInfo>(pVirtualHostInfoList->Get(i)))
         continue;

      // initialize SSL context for each SSLVirtualHostInfo
      VirtualHostInfo *pVHInfo = pVirtualHostInfoList->Get(i);
      for (int j = 0; j < static_cast<int>(pVHInfo->GetVHList().GetCount()); ++j) {
         SSLVHost * pSingleHost = static_cast<SSLVHost*>(pVHInfo->GetVHList().Get(j));
         SSLContext *pSSLCtxt = new SSLContext(true, rootCertPath, pSingleHost->needClientAuth);
         pSSLCtxt->SetCertChain(*pSingleHost->serverCertPath, pSingleHost->pPrivateKeyPwdCb);
         if (j == DEFAULT_SSLCONTEXT_INDEX) {
            ISSLFindMatchingHostNameCallback *pFindHostNameCb = 
                          new ServerFindMatchingHostNameCallback(const_cast<ListCppObjT<BaseObject>*>(&pVHInfo->GetVHList()));
            pSSLCtxt->SetServerNameCallback(pFindHostNameCb);
         }
         pSingleHost->SetAttachingSSLContext(pSSLCtxt);
      }
   }

}

SSLServerSocketFactory::~SSLServerSocketFactory(void)
{
}


/**
 *  This can only create non-SSL client socket
 */
Socket * SSLServerSocketFactory::CreateSocket(bool isSSL, const CFString &serverAddress, u16_t listenPort)
{
   return SocketFactory::CreateSocket(isSSL, serverAddress, listenPort);
}

/**
 *  @return server socket created
 *  @return NULL if error
 */
PRIVATE Socket * __CreateSSLSocket(SSLVHost &singleHost, const CFString &bindAddr, u16_t listenPort)
{
   SSLServerSocket *pServSock = singleHost.pSSLCtxt->CreateServerSocket(singleHost.pPeerCertCb);
   try {
      pServSock->Create(bindAddr, listenPort);
   }
   catch(SSLSocketException &e) {
      e.GetErrCode();
      cppmempool::Release(pServSock);
      pServSock = NULL;
   }

   return pServSock;
}

/**
 *  Create listening sockets for a SSL server
 */
bool SSLServerSocketFactory::CreateListeningSockets(collections::ListCppObjT<Socket> &sockList, int backLog)
{
   if (!pVirtualHostInfoList)
      return false;

   for (int i = 0; i < static_cast<int>(pVirtualHostInfoList->GetCount()); ++i) {
      VirtualHostInfo *pVHInfo = pVirtualHostInfoList->Get(i);
      // If SSL, use the first SSL context to create socket but may be changed later on
      // See SSLContext::SetServerNameCallback
      Socket *pListenSock = cppflib::InstanceOf<SSLVirtualHostInfo>(pVHInfo) ? 
                              __CreateSSLSocket(*static_cast<SSLVHost*>(pVHInfo->GetVHList().Get(DEFAULT_SSLCONTEXT_INDEX)), 
                                                pVHInfo->GetBindAddr(), pVHInfo->GetListenPort()) :
                              CreateSocket(false, pVHInfo->GetBindAddr(), pVHInfo->GetListenPort());
      
      if (!pListenSock) {
         sockList.Clear();
         break;
      }
      pListenSock->Listen(backLog);
      sockList.AddTail(pListenSock);
   }

   return (sockList.GetCount() > 0);
}

} // end of namespace ssl

} // end of namespace net

} // end of namespace cppflib
